#!usr/bin/env python
# encoding:utf-8
from __future__ import division
from PIL import Image

'''
功能： mat数据解析处理模块，实现mat格式数据转化为json格式数据
'''

import sys
import csv
import json
from scipy.io import loadmat


import os


def get_pic_file(dir_path):
    file_list = []
    r = ''
    for root, dirs, files in os.walk(dir_path):
        for file in files:
            if file.endswith("jpg") or file.endswith("png"):
                file_list.append(file)
                # file_list.append(os.path.join(root, file))

        r = root
    file_list.sort(key = lambda x:int(x.split('.')[0]))

    return [os.path.join(r, file) for file in file_list]


def get_pic_size(file_list):
    size_list = []
    for file in file_list:
        img = Image.open(file)
        size_list.append(img.size)
    return size_list





def getx1y1x2y2list(data='a.mat'):
    '''
    将.mat文件转化为json文件
    '''
    mat = loadmat(data)
    print(mat)
    labels = mat['box'].tolist()
    for label in labels:
        label[2] = label[0]+label[2]
        label[3] = label[1]+label[3]

    return labels
    # label = [one[0] for one in labels]
    # # datas = mat['data'].tolist()
    # res_list = []
    # for i in range(len(labels)):
    #     # one_data = datas[i]
    #     one_label = label[i]
    #     # one_data.append(one_label)
    #     res_list.append(one_label)
    # with open(save_path, 'w') as f:
    #     f.write(json.dumps(res_list))

# 将x1, y1, x2, y2转换成yolov5所需要的x, y, w, h格式
def xyxy2xywh(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[2]) / 2 * dw
    y = (box[1] + box[3]) / 2 * dh
    w = (box[2] - box[0]) * dw
    h = (box[3] - box[1]) * dh
    return (x, y, w, h)  # 返回的都是标准化后的值

def getwhlist(size_list, box_list):
    wh_list=[]
    for i in range(len(size_list)):
        wh_list.append(xyxy2xywh(size_list[i],box_list[i]))

    return wh_list

if __name__ == '__main__':
    str_ = "26"
    path = '''D:\\Program\\yolov5\\Yolov5_StrongSORT_OSNet\\yolov5\\datasets\\uav\\set4\\online_drone_dataset\\image&label-20220708T015451Z-001\\image_label\\'''+str_
    #path = '''D:\\Program\\yolov5\\Yolov5_StrongSORT_OSNet\\yolov5\\datasets\\uav\\set4\\usr_drone_001\\usc drone-001\\Drone8\\img'''
    file_list = get_pic_file(path)
    size_list = get_pic_size(file_list)
    box_list = getx1y1x2y2list('D:\\Program\\yolov5\\Yolov5_StrongSORT_OSNet\\yolov5\\datasets\\uav\\set4\\online_drone_dataset\\image&label-20220708T015451Z-001\\image_label\\'+str_+'\\anotation.mat')
    #box_list = getx1y1x2y2list('D:\\Program\\yolov5\\Yolov5_StrongSORT_OSNet\\yolov5\\datasets\\uav\\set4\\usr_drone_001\\usc drone-001\\Drone8\\img\\ground-truth.mat')
    print(file_list)
    print(box_list)
    wh_list = getwhlist(size_list,box_list)
    print(wh_list)



    for i in range(len(file_list)):
        with open(file_list[i].split('.')[0]+'.txt','w') as f:
            f.write("0 " + ' '.join([str(item) for item in wh_list[i]]))


    # dataParse2Json('anotation.mat')